// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
/**
 * This file tests the performance of each JsonXXXFunction.call() with
 * expression framework and Pollux vectors.
 */
#include <melon/benchmark.h>
#include <melon/init/init.h>
#include <pollux/functions/registerer.h>
#include <pollux/functions/lib/benchmarks/FunctionBenchmarkBase.h>
#include <pollux/functions/prestosql/json_functions.h>
#include <pollux/functions/prestosql/json/json_extractor.h>
#include <pollux/functions/prestosql/types/json_registration.h>
#include <pollux/functions/prestosql/types/json_type.h>

namespace kumo::pollux::functions {
void registerJsonVectorFunctions() {
  POLLUX_REGISTER_VECTOR_FUNCTION(udf_json_extract, "json_extract");
}
} // namespace kumo::pollux::functions

namespace kumo::pollux::functions::prestosql {
namespace {

template <typename T>
struct FollyIsJsonScalarFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  MELON_ALWAYS_INLINE void call(bool& result, const arg_type<Json>& json) {
    auto parsedJson = melon::parseJson(json);
    result = parsedJson.isNumber() || parsedJson.isString() ||
        parsedJson.isBool() || parsedJson.isNull();
  }
};

// jsonExtractScalar(json, json_path) -> varchar
// Current implementation support UTF-8 in json, but not in json_path.
// Like jsonExtract(), but returns the result value as a string (as opposed
// to being encoded as JSON). The value referenced by json_path must be a scalar
// (boolean, number or string)
template <typename T>
struct FollyJsonExtractScalarFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  MELON_ALWAYS_INLINE bool call(
      out_type<Varchar>& result,
      const arg_type<Json>& json,
      const arg_type<Varchar>& jsonPath) {
    const melon::StringPiece& jsonStringPiece = json;
    const melon::StringPiece& jsonPathStringPiece = jsonPath;
    auto extractResult =
        jsonExtractScalar(jsonStringPiece, jsonPathStringPiece);
    if (extractResult.hasValue()) {
      UDFOutputString::assign(result, *extractResult);
      return true;
    } else {
      return false;
    }
  }
};

template <typename T>
struct FollyJsonExtractFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  MELON_ALWAYS_INLINE bool call(
      out_type<Json>& result,
      const arg_type<Json>& json,
      const arg_type<Varchar>& jsonPath) {
    auto extractResult =
        jsonExtract(melon::StringPiece(json), melon::StringPiece(jsonPath));
    if (!extractResult.hasValue() || extractResult.value().isNull()) {
      return false;
    }

    melon::json::serialization_opts opts;
    opts.sort_keys = true;
    melon::json::serialize(*extractResult, opts);

    UDFOutputString::assign(
        result, melon::json::serialize(*extractResult, opts));
    return true;
  }
};

template <typename T>
struct FollyJsonArrayLengthFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  MELON_ALWAYS_INLINE bool call(int64_t& result, const arg_type<Json>& json) {
    auto parsedJson = melon::parseJson(json);
    if (!parsedJson.is_array()) {
      return false;
    }

    result = parsedJson.size();
    return true;
  }
};

template <typename T>
struct FollyJsonArrayContainsFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  template <typename TInput>
  MELON_ALWAYS_INLINE bool
  call(bool& result, const arg_type<Json>& json, const TInput& value) {
    auto parsedJson = melon::parseJson(json);
    if (!parsedJson.is_array()) {
      return false;
    }

    result = false;
    for (const auto& v : parsedJson) {
      if constexpr (std::is_same_v<TInput, bool>) {
        if (v.isBool() && v == value) {
          result = true;
          break;
        }
      } else if constexpr (std::is_same_v<TInput, int64_t>) {
        if (v.isInt() && v == value) {
          result = true;
          break;
        }
      } else if constexpr (std::is_same_v<TInput, double>) {
        if (v.is_double() && v == value) {
          result = true;
          break;
        }
      } else {
        if (v.isString() && v == value) {
          result = true;
          break;
        }
      }
    }
    return true;
  }
};

template <typename T>
struct FollyJsonSizeFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  MELON_ALWAYS_INLINE bool call(
      int64_t& result,
      const arg_type<Json>& json,
      const arg_type<Varchar>& jsonPath) {
    const melon::StringPiece& jsonStringPiece = json;
    const melon::StringPiece& jsonPathStringPiece = jsonPath;
    auto extractResult = jsonExtract(jsonStringPiece, jsonPathStringPiece);
    if (!extractResult.has_value()) {
      return false;
    }
    // The size of the object or array is the number of members, otherwise the
    // size is zero
    if (extractResult->is_array() || extractResult->isObject()) {
      result = extractResult->size();
    } else {
      result = 0;
    }

    return true;
  }
};

const std::string smallJson = R"({"k1":"v1"})";

class JsonBenchmark : public pollux::functions::test::FunctionBenchmarkBase {
 public:
  JsonBenchmark() : FunctionBenchmarkBase() {
    registerJsonType();
    register_function<IsJsonScalarFunction, bool, Json>({"is_json_scalar"});
    register_function<FollyIsJsonScalarFunction, bool, Json>(
        {"melon_is_json_scalar"});
    register_function<JsonArrayContainsFunction, bool, Json, bool>(
        {"json_array_contains"});
    register_function<FollyJsonArrayContainsFunction, bool, Json, bool>(
        {"melon_json_array_contains"});
    register_function<JsonArrayLengthFunction, int64_t, Json>(
        {"json_array_length"});
    register_function<FollyJsonArrayLengthFunction, int64_t, Json>(
        {"melon_json_array_length"});
    register_function<FollyJsonExtractScalarFunction, Varchar, Json, Varchar>(
        {"melon_json_extract_scalar"});
    register_function<JsonExtractScalarFunction, Varchar, Json, Varchar>(
        {"json_extract_scalar"});
    register_function<FollyJsonExtractFunction, Varchar, Json, Varchar>(
        {"melon_json_extract"});
    register_function<JsonSizeFunction, int64_t, Json, Varchar>({"json_size"});
    register_function<FollyJsonSizeFunction, int64_t, Json, Varchar>(
        {"melon_json_size"});
    registerJsonVectorFunctions();
  }

  std::string prepareData(int jsonSize) {
    std::string jsonData = R"({"key": [)" + smallJson;
    for (int i = 0; i < jsonSize / 10; i++) {
      jsonData += "," + smallJson;
    }
    jsonData += "]}";
    return jsonData;
  }

  pollux::VectorPtr makeJsonData(const std::string& json, int vectorSize) {
    auto jsonVector =
        vectorMaker_.flat_vector<pollux::StringView>(vectorSize, JSON());
    for (auto i = 0; i < vectorSize; i++) {
      jsonVector->set(i, pollux::StringView(json));
    }
    return jsonVector;
  }

  void runWithJson(
      int iter,
      int vectorSize,
      const std::string& fnName,
      const std::string& json) {
    melon::BenchmarkSuspender suspender;

    auto jsonVector = makeJsonData(json, vectorSize);

    auto row_vector = vectorMaker_.row_vector({jsonVector});
    auto exprSet =
        compileExpression(fmt::format("{}(c0)", fnName), row_vector->type());
    suspender.dismiss();
    doRun(iter, exprSet, row_vector);
  }

  void runWithJsonExtract(
      int iter,
      int vectorSize,
      const std::string& fnName,
      const std::string& json,
      const std::string& path) {
    melon::BenchmarkSuspender suspender;

    auto jsonVector = makeJsonData(json, vectorSize);
    auto pathVector = vectorMaker_.constant_vector(
        std::vector<std::optional<StringView>>(vectorSize, StringView(path)));

    auto row_vector = vectorMaker_.row_vector({jsonVector, pathVector});
    auto exprSet =
        compileExpression(fmt::format("{}(c0, c1)", fnName), row_vector->type());
    suspender.dismiss();
    doRun(iter, exprSet, row_vector);
  }

  void runWithJsonContains(
      int iter,
      int vectorSize,
      const std::string& fnName,
      const std::string& json) {
    melon::BenchmarkSuspender suspender;

    auto jsonVector = makeJsonData(json, vectorSize);
    auto boolVector = vectorMaker_.flat_vector<bool>({true});

    auto row_vector = vectorMaker_.row_vector({jsonVector, boolVector});
    auto exprSet =
        compileExpression(fmt::format("{}(c0, c1)", fnName), row_vector->type());
    suspender.dismiss();
    doRun(iter, exprSet, row_vector);
  }

  void doRun(
      const int iter,
      pollux::exec::ExprSet& exprSet,
      const pollux::RowVectorPtr& row_vector) {
    uint32_t cnt = 0;
    for (auto i = 0; i < iter; i++) {
      cnt += evaluate(exprSet, row_vector)->size();
    }
    melon::doNotOptimizeAway(cnt);
  }
};

void FollyIsJsonScalar(int iter, int vectorSize, int jsonSize) {
  melon::BenchmarkSuspender suspender;
  JsonBenchmark benchmark;
  auto json = benchmark.prepareData(jsonSize);
  suspender.dismiss();
  benchmark.runWithJson(iter, vectorSize, "melon_is_json_scalar", json);
}

void SIMDIsJsonScalar(int iter, int vectorSize, int jsonSize) {
  melon::BenchmarkSuspender suspender;
  JsonBenchmark benchmark;
  auto json = benchmark.prepareData(jsonSize);
  suspender.dismiss();
  benchmark.runWithJson(iter, vectorSize, "is_json_scalar", json);
}

void FollyJsonArrayContains(int iter, int vectorSize, int jsonSize) {
  melon::BenchmarkSuspender suspender;
  JsonBenchmark benchmark;
  auto json = benchmark.prepareData(jsonSize);
  suspender.dismiss();
  benchmark.runWithJsonContains(
      iter, vectorSize, "melon_json_array_contains", json);
}

void SIMDJsonArrayContains(int iter, int vectorSize, int jsonSize) {
  melon::BenchmarkSuspender suspender;
  JsonBenchmark benchmark;
  auto json = benchmark.prepareData(jsonSize);
  suspender.dismiss();
  benchmark.runWithJsonContains(iter, vectorSize, "json_array_contains", json);
}

void FollyJsonArrayLength(int iter, int vectorSize, int jsonSize) {
  melon::BenchmarkSuspender suspender;
  JsonBenchmark benchmark;
  auto json = benchmark.prepareData(jsonSize);
  suspender.dismiss();
  benchmark.runWithJson(iter, vectorSize, "melon_json_array_length", json);
}

void SIMDJsonArrayLength(int iter, int vectorSize, int jsonSize) {
  melon::BenchmarkSuspender suspender;
  JsonBenchmark benchmark;
  auto json = benchmark.prepareData(jsonSize);
  suspender.dismiss();
  benchmark.runWithJson(iter, vectorSize, "json_array_length", json);
}

void FollyJsonExtractScalar(int iter, int vectorSize, int jsonSize) {
  melon::BenchmarkSuspender suspender;
  JsonBenchmark benchmark;
  auto json = benchmark.prepareData(jsonSize);
  suspender.dismiss();
  benchmark.runWithJsonExtract(
      iter, vectorSize, "melon_json_extract_scalar", json, "$.key[7].k1");
}

void SIMDJsonExtractScalar(int iter, int vectorSize, int jsonSize) {
  melon::BenchmarkSuspender suspender;
  JsonBenchmark benchmark;
  auto json = benchmark.prepareData(jsonSize);
  suspender.dismiss();
  benchmark.runWithJsonExtract(
      iter, vectorSize, "json_extract_scalar", json, "$.key[7].k1");
}

void FollyJsonExtract(int iter, int vectorSize, int jsonSize) {
  melon::BenchmarkSuspender suspender;
  JsonBenchmark benchmark;
  auto json = benchmark.prepareData(jsonSize);
  suspender.dismiss();
  benchmark.runWithJsonExtract(
      iter, vectorSize, "melon_json_extract", json, "$.key[*].k1");
}

void SIMDJsonExtract(int iter, int vectorSize, int jsonSize) {
  melon::BenchmarkSuspender suspender;
  JsonBenchmark benchmark;
  auto json = benchmark.prepareData(jsonSize);
  suspender.dismiss();
  benchmark.runWithJsonExtract(
      iter, vectorSize, "json_extract", json, "$.key[*].k1");
}

void FollyJsonSize(int iter, int vectorSize, int jsonSize) {
  melon::BenchmarkSuspender suspender;
  JsonBenchmark benchmark;
  auto json = benchmark.prepareData(jsonSize);
  suspender.dismiss();
  benchmark.runWithJsonExtract(
      iter, vectorSize, "melon_json_size", json, "$.key");
}

void SIMDJsonSize(int iter, int vectorSize, int jsonSize) {
  melon::BenchmarkSuspender suspender;
  JsonBenchmark benchmark;
  auto json = benchmark.prepareData(jsonSize);
  suspender.dismiss();
  benchmark.runWithJsonExtract(iter, vectorSize, "json_size", json, "$.key");
}

BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(FollyIsJsonScalar, 100_iters_10bytes_size, 100, 10);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDIsJsonScalar,
    100_iters_10bytes_size,
    100,
    10);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(FollyIsJsonScalar, 100_iters_100bytes_size, 100, 100);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDIsJsonScalar,
    100_iters_100bytes_size,
    100,
    100);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(FollyIsJsonScalar, 100_iters_1000bytes_size, 100, 1000);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDIsJsonScalar,
    100_iters_1000bytes_size,
    100,
    1000);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(FollyIsJsonScalar, 100_iters_10000bytes_size, 100, 10000);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDIsJsonScalar,
    100_iters_10000bytes_size,
    100,
    10000);
BENCHMARK_DRAW_LINE();

BENCHMARK_DRAW_LINE();
BENCHMARK_NAMED_PARAM(FollyJsonArrayContains, 100_iters_10bytes_size, 100, 10);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonArrayContains,
    100_iters_10bytes_size,
    100,
    10);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(
    FollyJsonArrayContains,
    100_iters_100bytes_size,
    100,
    100);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonArrayContains,
    100_iters_100bytes_size,
    100,
    100);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(
    FollyJsonArrayContains,
    100_iters_1000bytes_size,
    100,
    1000);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonArrayContains,
    100_iters_1000bytes_size,
    100,
    1000);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(
    FollyJsonArrayContains,
    100_iters_10000bytes_size,
    100,
    10000);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonArrayContains,
    100_iters_10000bytes_size,
    100,
    10000);
BENCHMARK_DRAW_LINE();

BENCHMARK_DRAW_LINE();
BENCHMARK_NAMED_PARAM(FollyJsonArrayLength, 100_iters_10bytes_size, 100, 10);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonArrayLength,
    100_iters_10bytes_size,
    100,
    10);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(FollyJsonArrayLength, 100_iters_100bytes_size, 100, 100);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonArrayLength,
    100_iters_100bytes_size,
    100,
    100);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(
    FollyJsonArrayLength,
    100_iters_1000bytes_size,
    100,
    1000);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonArrayLength,
    100_iters_1000bytes_size,
    100,
    1000);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(
    FollyJsonArrayLength,
    100_iters_10000bytes_size,
    100,
    10000);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonArrayLength,
    100_iters_10000bytes_size,
    100,
    10000);
BENCHMARK_DRAW_LINE();

BENCHMARK_DRAW_LINE();
BENCHMARK_NAMED_PARAM(FollyJsonExtractScalar, 100_iters_10bytes_size, 100, 10);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonExtractScalar,
    100_iters_10bytes_size,
    100,
    10);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(
    FollyJsonExtractScalar,
    100_iters_100bytes_size,
    100,
    100);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonExtractScalar,
    100_iters_100bytes_size,
    100,
    100);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(
    FollyJsonExtractScalar,
    100_iters_1000bytes_size,
    100,
    1000);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonExtractScalar,
    100_iters_1000bytes_size,
    100,
    1000);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(
    FollyJsonExtractScalar,
    100_iters_10000bytes_size,
    100,
    10000);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonExtractScalar,
    100_iters_10000bytes_size,
    100,
    10000);
BENCHMARK_DRAW_LINE();

BENCHMARK_DRAW_LINE();
BENCHMARK_NAMED_PARAM(FollyJsonExtract, 100_iters_10bytes_size, 100, 10);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonExtract,
    100_iters_10bytes_size,
    100,
    10);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(FollyJsonExtract, 100_iters_100bytes_size, 100, 100);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonExtract,
    100_iters_100bytes_size,
    100,
    100);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(FollyJsonExtract, 100_iters_1000bytes_size, 100, 1000);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonExtract,
    100_iters_1000bytes_size,
    100,
    1000);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(FollyJsonExtract, 100_iters_10000bytes_size, 100, 10000);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonExtract,
    100_iters_10000bytes_size,
    100,
    10000);
BENCHMARK_DRAW_LINE();

BENCHMARK_DRAW_LINE();
BENCHMARK_NAMED_PARAM(FollyJsonSize, 100_iters_10bytes_size, 100, 10);
BENCHMARK_RELATIVE_NAMED_PARAM(SIMDJsonSize, 100_iters_10bytes_size, 100, 10);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(FollyJsonSize, 100_iters_100bytes_size, 100, 100);
BENCHMARK_RELATIVE_NAMED_PARAM(SIMDJsonSize, 100_iters_100bytes_size, 100, 100);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(FollyJsonSize, 100_iters_1000bytes_size, 100, 1000);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonSize,
    100_iters_1000bytes_size,
    100,
    1000);
BENCHMARK_DRAW_LINE();

BENCHMARK_NAMED_PARAM(FollyJsonSize, 100_iters_10000bytes_size, 100, 10000);
BENCHMARK_RELATIVE_NAMED_PARAM(
    SIMDJsonSize,
    100_iters_10000bytes_size,
    100,
    10000);
BENCHMARK_DRAW_LINE();

} // namespace
} // namespace kumo::pollux::functions::prestosql

int main(int argc, char** argv) {
  melon::Init init{&argc, &argv};

  kumo::pollux::memory::MemoryManager::initialize({});

  melon::runBenchmarks();
  return 0;
}
